#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""带跟踪功能的YOLO检测引擎

集成了目标检测和跟踪功能的高性能引擎
"""

import cv2
import numpy as np
import logging
import time
from typing import List, Dict, Optional, Tuple, Any
from collections import defaultdict

from core.detection.yolo_engine import YOLODetectionEngine
from config import get_config

# 设置日志记录器
logger = logging.getLogger(__name__)

class YOLOTrackingEngine(YOLODetectionEngine):
    """带跟踪功能的YOLO检测引擎"""
    
    def __init__(self, model_path: Optional[str] = None, device: str = 'auto', **kwargs):
        """初始化带跟踪功能的YOLO检测引擎
        
        Args:
            model_path: 模型文件路径
            device: 运行设备，'auto'、'cuda'、'cpu'
            **kwargs: 其他配置参数
                - tracking_config: 跟踪器配置参数
        """
        super().__init__(model_path=model_path, device=device, **kwargs)
        
        # 跟踪相关属性
        self.track_history = defaultdict(dict)
        self.track_ids = 0
        self.last_track_time = time.time()
        
        # 跟踪器配置
        self.tracking_config = kwargs.get('tracking_config', {})
        self.tracking_config.setdefault('max_age', 30)  # 最大未更新时间（帧）
        self.tracking_config.setdefault('iou_threshold', 0.3)  # IoU匹配阈值
        self.tracking_config.setdefault('max_trajectory_length', 100)  # 最大轨迹长度
        
        # 速度和方向计算配置
        self.tracking_config.setdefault('speed_calculation_enabled', True)  # 是否启用速度计算
        
        logger.info("YOLOTrackingEngine初始化完成")
    
    def track(self, detections: List[Dict[str, Any]], frame: np.ndarray) -> List[Dict[str, Any]]:
        """跟踪检测到的目标
        
        Args:
            detections: 检测结果列表
            frame: 当前图像帧
            
        Returns:
            List[Dict]: 跟踪结果列表，包含跟踪ID和轨迹信息
        """
        if not detections:
            return []
        
        try:
            current_time = time.time()
            tracked_objects = []
            used_detection_indices = set()
            new_track_history = defaultdict(dict)
            
            # 遍历现有跟踪器，尝试匹配新检测
            for track_id, track_info in self.track_history.items():
                best_match_idx = -1
                best_iou = 0
                
                # 查找最佳匹配的检测
                for i, detection in enumerate(detections):
                    if i in used_detection_indices:
                        continue
                    
                    iou = self._calculate_iou(track_info['last_bbox'], detection['bbox'])
                    
                    # 检查是否是同一类别
                    if track_info['class_id'] != detection['class_id']:
                        continue
                    
                    if iou > best_iou and iou >= self.tracking_config['iou_threshold']:
                        best_iou = iou
                        best_match_idx = i
                
                # 如果找到匹配项
                if best_match_idx >= 0:
                    detection = detections[best_match_idx]
                    
                    # 更新跟踪信息
                    track_info['last_bbox'] = detection['bbox']
                    track_info['last_center'] = detection['center']
                    track_info['last_time'] = current_time
                    track_info['confidence'] = detection['confidence']
                    
                    # 更新轨迹历史
                    track_info['trajectory'].append(detection['center'])
                    if len(track_info['trajectory']) > self.tracking_config['max_trajectory_length']:
                        track_info['trajectory'].pop(0)
                    
                    # 计算移动方向和速度
                    if self.tracking_config['speed_calculation_enabled'] and len(track_info['trajectory']) > 1:
                        prev_center = track_info['trajectory'][-2]
                        current_center = track_info['trajectory'][-1]
                        track_info['direction'] = (
                            current_center[0] - prev_center[0],
                            current_center[1] - prev_center[1]
                        )
                        
                        time_diff = current_time - track_info['times'][-1]
                        if time_diff > 0:
                            track_info['speed'] = np.sqrt(
                                (track_info['direction'][0] / time_diff)**2 +
                                (track_info['direction'][1] / time_diff)** 2
                            )
                    
                    track_info['times'].append(current_time)
                    if len(track_info['times']) > self.tracking_config['max_trajectory_length']:
                        track_info['times'].pop(0)
                    
                    new_track_history[track_id] = track_info
                    used_detection_indices.add(best_match_idx)
                    
                    # 添加到跟踪结果
                    tracked_objects.append({
                        'track_id': track_id,
                        **detection,
                        'trajectory': track_info['trajectory'],
                        'direction': track_info.get('direction', (0, 0)),
                        'speed': track_info.get('speed', 0),
                        'tracking_time': current_time - track_info['start_time'],
                        'tracked_frames': len(track_info['times'])
                    })
            
            # 为未匹配的检测创建新的跟踪器
            max_age = self.tracking_config['max_age']
            for i, detection in enumerate(detections):
                if i not in used_detection_indices:
                    # 创建新的跟踪ID
                    self.track_ids += 1
                    track_id = self.track_ids
                    
                    # 初始化跟踪信息
                    track_info = {
                        'last_bbox': detection['bbox'],
                        'last_center': detection['center'],
                        'last_time': current_time,
                        'start_time': current_time,
                        'class_id': detection['class_id'],
                        'class_name': detection['class_name'],
                        'confidence': detection['confidence'],
                        'trajectory': [detection['center']],
                        'times': [current_time],
                        'direction': (0, 0),
                        'speed': 0
                    }
                    
                    new_track_history[track_id] = track_info
                    
                    # 添加到跟踪结果
                    tracked_objects.append({
                        'track_id': track_id,
                        **detection,
                        'trajectory': track_info['trajectory'],
                        'direction': track_info['direction'],
                        'speed': track_info['speed'],
                        'tracking_time': 0,
                        'tracked_frames': 1
                    })
            
            # 更新跟踪历史
            self.track_history = new_track_history
            
            # 定期清理长时间未更新的跟踪器
            if current_time - self.last_track_time > 1.0:  # 每秒清理一次
                self._cleanup_old_tracks(current_time, max_age)
                self.last_track_time = current_time
            
            return tracked_objects
            
        except Exception as e:
            logger.error(f"跟踪失败: {str(e)}")
            return detections  # 失败时返回原始检测结果
    
    def _calculate_iou(self, box1: List[int], box2: List[int]) -> float:
        """计算两个边界框的IoU
        
        Args:
            box1: [x1, y1, x2, y2]
            box2: [x1, y1, x2, y2]
            
        Returns:
            float: IoU值
        """
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])
        
        # 计算交集面积
        intersection = max(0, x2 - x1) * max(0, y2 - y1)
        
        # 计算并集面积
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        union = area1 + area2 - intersection
        
        if union == 0:
            return 0
        
        return intersection / union
    
    def _cleanup_old_tracks(self, current_time: float, max_age: float):
        """清理过期的跟踪器
        
        Args:
            current_time: 当前时间
            max_age: 最大允许的未更新时间（秒）
        """
        expired_ids = []
        for track_id, track_info in self.track_history.items():
            if current_time - track_info['last_time'] > max_age:
                expired_ids.append(track_id)
        
        for track_id in expired_ids:
            del self.track_history[track_id]
        
        if expired_ids:
            logger.debug(f"清理了 {len(expired_ids)} 个过期跟踪器")
    
    def get_tracking_stats(self) -> Dict[str, Any]:
        """获取跟踪统计信息
        
        Returns:
            Dict: 跟踪统计信息
        """
        stats = {
            'active_tracks': len(self.track_history),
            'track_ids_count': self.track_ids,
            'track_history_size': sum(len(v['trajectory']) for v in self.track_history.values()),
            'max_trajectory_length': self.tracking_config['max_trajectory_length'],
            'max_age': self.tracking_config['max_age']
        }
        return stats
    
    def visualize_results(self, frame: np.ndarray, detections: List[Dict[str, Any]], 
                         show_confidence: bool = True, show_tracks: bool = True, 
                         tracked_objects: Optional[List[Dict[str, Any]]] = None) -> np.ndarray:
        """可视化检测和跟踪结果
        
        Args:
            frame: 输入图像帧
            detections: 检测结果列表
            show_confidence: 是否显示置信度
            show_tracks: 是否显示跟踪轨迹
            tracked_objects: 跟踪结果列表
            
        Returns:
            np.ndarray: 可视化后的图像
        """
        # 先调用父类的可视化方法绘制检测框
        result_frame = super().visualize_results(
            frame, detections, show_confidence=show_confidence,
            show_tracks=False, tracked_objects=None
        )
        
        # 如果需要显示跟踪信息且有跟踪结果
        if show_tracks and tracked_objects:
            # 为不同类别定义颜色
            colors = {
                0: (255, 0, 0),  # person - 蓝色
                1: (0, 255, 0),  # bicycle - 绿色
                2: (0, 0, 255),  # car - 红色
                3: (255, 255, 0),  # motorbike - 青色
                5: (255, 0, 255),  # bus - 紫色
                7: (0, 255, 255),  # truck - 黄色
                'default': (128, 128, 128)  # 默认 - 灰色
            }
            
            for track_obj in tracked_objects:
                track_id = track_obj['track_id']
                bbox = track_obj['bbox']
                class_id = track_obj['class_id']
                confidence = track_obj['confidence']
                trajectory = track_obj.get('trajectory', [])
                speed = track_obj.get('speed', 0)
                
                # 获取颜色
                color = colors.get(class_id, colors['default'])
                
                # 绘制边界框（用虚线表示跟踪对象）
                x1, y1, x2, y2 = bbox
                cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 2)
                
                # 绘制跟踪ID
                track_text = f"ID: {track_id}"
                cv2.putText(
                    result_frame,
                    track_text,
                    (x1 + 5, y1 - 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    color,
                    2
                )
                
                # 绘制速度信息（如果可用）
                if speed > 0:
                    speed_text = f"Speed: {speed:.1f}"
                    cv2.putText(
                        result_frame,
                        speed_text,
                        (x1 + 5, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        color,
                        2
                    )
                
                # 绘制轨迹线
                if len(trajectory) > 1:
                    for i in range(1, len(trajectory)):
                        cv2.line(
                            result_frame,
                            (int(trajectory[i-1][0]), int(trajectory[i-1][1])),
                            (int(trajectory[i][0]), int(trajectory[i][1])),
                            color,
                            2
                        )
        
        return result_frame
    
    def get_performance_metrics(self) -> Dict[str, Any]:
        """获取检测引擎性能指标，包含跟踪统计
        
        Returns:
            Dict: 性能指标字典
        """
        metrics = super().get_performance_metrics()
        metrics.update(self.get_tracking_stats())
        return metrics
    
    def reset_tracking(self):
        """重置跟踪状态"""
        self.track_history.clear()
        self.track_ids = 0
        self.last_track_time = time.time()
        logger.info("跟踪状态已重置")
    
    def set_tracking_config(self, config: Dict[str, Any]):
        """设置跟踪器配置
        
        Args:
            config: 跟踪器配置字典
        """
        self.tracking_config.update(config)
        logger.info(f"跟踪配置已更新: {self.tracking_config}")
        
        # 如果更新了最大轨迹长度，截断现有轨迹
        if 'max_trajectory_length' in config:
            for track_info in self.track_history.values():
                if len(track_info['trajectory']) > self.tracking_config['max_trajectory_length']:
                    track_info['trajectory'] = track_info['trajectory'][-self.tracking_config['max_trajectory_length']:]
                    track_info['times'] = track_info['times'][-self.tracking_config['max_trajectory_length']:]
    
    def release(self):
        """释放检测引擎资源"""
        # 清理跟踪历史
        self.reset_tracking()
        # 调用父类的释放方法
        super().release()